import numpy

from histogram import HistogramCollection, Histogram1D, Histogram2D

hcoll = HistogramCollection(
    [Histogram1D(name='rs_dvf_m_pull',
                 nbins=100,
                 arange=(-5, 5),
                 xlabel=r'$M(\left[K\pi\right]_{RS})$ Pull',
                 xunits=r'',
                 assoc_cols='rs_dvf_m_pull'),
     Histogram1D(name='ws_dvf_m_pull',
                 nbins=100,
                 arange=(-5, 5),
                 xlabel=r'$M(\left[K\pi\right]_{WS})$ Pull',
                 xunits=r'',
                 assoc_cols='ws_dvf_m_pull'),
     Histogram1D(name='rs_dvf_m_unc',
                 nbins=100,
                 arange=(0, 1e-2),
                 xlabel=r'$M(\left[K\pi\right]_{RS})$ Unc',
                 xunits=r'GeV/c',
                 assoc_cols='rs_dvf_m_unc'),
     Histogram1D(name='ws_dvf_m_unc',
                 nbins=100,
                 arange=(0, 1e-2),
                 xlabel=r'$M(\left[K\pi\right]_{WS})$ Unc',
                 xunits=r'GeV/c',
                 assoc_cols='ws_dvf_m_unc'),
     Histogram1D(name='rs_dvf_cl',
                 nbins=100,
                 arange=(0, 1),
                 xlabel=r'$\left[K\pi\right]_{RS}$ vertex fit confidence',
                 xunits=r'',
                 assoc_cols='rs_dvf_cl'),
     Histogram1D(name='ws_dvf_cl',
                 nbins=100,
                 arange=(0, 1),
                 xlabel=r'$\left[K\pi\right]_{WS}$ vertex fit confidence',
                 xunits=r'',
                 assoc_cols='ws_dvf_cl'),
     Histogram1D(name='rs_dvf_chi2',
                 nbins=100,
                 arange=(0, 100),
                 xlabel=r'$\left[K\pi\right]_{RS}$ vertex fit $\chi^2$',
                 xunits=r'',
                 assoc_cols='rs_dvf_chi2'),
     Histogram1D(name='ws_dvf_chi2',
                 nbins=100,
                 arange=(0, 100),
                 xlabel=r'$\left[K\pi\right]_{WS}$ vertex fit $\chi^2$',
                 xunits=r'',
                 assoc_cols='ws_dvf_chi2'),
     Histogram1D(name='m_pipipi',
                 nbins=1010,
                 arange=(0.4, 5.5),
                 xlabel=r'M($\pi_s\pi\pi$)',
                 xunits=r'GeV',
                 assoc_cols='m_pipipi'),
     Histogram1D(name='m_kkk',
                 nbins=1000,
                 arange=(2.3, 3.3),
                 xlabel=r'M($K_sKK$)',
                 xunits=r'GeV',
                 assoc_cols='m_kkk'),
     Histogram1D(name='m_dstar_rs',
                 nbins=700,
                 arange=(1.3, 2.7),
                 xlabel=r'M($\pi_s + \left[K\pi\right]_{RS})$)',
                 xunits=r'GeV',
                 assoc_cols='m_dstar_rs'),
     Histogram1D(name='m_dstar_ws',
                 nbins=700,
                 arange=(1.3, 2.7),
                 xlabel=r'M($\pi_s + \left[K\pi\right]_{WS})$)',
                 xunits=r'GeV',
                 assoc_cols='m_dstar_ws'),
     Histogram1D(name='m_rs',
                 nbins=250,
                 arange=(1.75, 2.0),
                 xlabel=r'M($\left[K\pi\right]_{WS}$)',
                 xunits=r'GeV',
                 assoc_cols='m_rs',
                 ),
     Histogram1D(name='m_ws',
                 nbins=250,
                 arange=(1.75, 2.0),
                 xlabel=r'M($\left[K\pi\right]_{WS}$)',
                 xunits=r'GeV',
                 assoc_cols='m_ws',
                 ),
     Histogram1D(name='q_ws',
                 nbins=200,
                 arange=(0,0.02),
                 xlabel=r'$M(\pi_s\left[K\pi\right]_{WS})' \
                     + r'-M(\left[K\pi\right]_{WS})-M_\pi$',
                 xunits='GeV',
                 title=r'$Q$ for $\left[K\pi\right]_{WS}$',
                 assoc_cols='q_ws',
                 ),
     Histogram1D(name='q_rs',
                 nbins=200,
                 arange=(0,0.02),
                 xlabel=r'$M(\pi_s\left[K\pi\right]_{RS})-' \
                     + r'M(\left[K\pi\right]_{RS})-M_\pi$',
                 xunits='GeV',
                 title=r'$Q$ for $\left[K\pi\right]_{RS}$',
                 assoc_cols='q_rs',
                 ),
     Histogram1D(name='pcm_dstar_rs',
                 nbins=1250,
                 arange=(1,7),
                 xlabel=r'$|p^*|(\pi_s + \left[K\pi\right]_{RS})$',
                 xunits='GeV',
                 assoc_cols='pcm_dstar_rs',
                 ),
     Histogram1D(name='pcm_dstar_ws',
                 nbins=1250,
                 arange=(1,7),
                 xlabel=r'$|p^*|(\pi_s + \left[K\pi\right]_{WS})$',
                 xunits='GeV',
                 assoc_cols='pcm_dstar_ws',
                 ),
     Histogram1D(name='ws_theta_star',
                 nbins=1000,
                 arange=(-1,1),
                 xlabel=r'$\theta^*$ ($\left[K\pi\right]_{WS}$ system)',
                 xunits='GeV',
                 assoc_cols='ws_theta_star',
                 ),
     Histogram1D(name='rs_theta_star',
                 nbins=1000,
                 arange=(-1,1),
                 xlabel=r'$\theta^*$ ($\left[K\pi\right]_{RS}$ system)',
                 xunits='GeV',
                 assoc_cols='rs_theta_star',
                 ),
     ])


def neg_sqrt(col):
    return numpy.where(col < 0, -numpy.sqrt(numpy.fabs(col)),
                       numpy.sqrt(numpy.fabs(col)))

hcoll.colfilters['rs_dvf_m_unc'] = neg_sqrt
hcoll.colfilters['ws_dvf_m_unc'] = neg_sqrt
